// queue.js
const treeKill = require('tree-kill');

class MessageQueue {
    constructor() {
        this.queues = {};
    }

    clear(taskID) {
        // console.log('Queue clear taskID = ', taskID);
        if (this.queues[taskID]) {
            this.queues[taskID] = [];
        }
    }

    // Add a message to the queue of the specific taskID
    enqueue(taskID, message, retainAfterPeek = true) {
        if (!this.queues[taskID]) {
            this.queues[taskID] = [];
        }

        // Add timestamp to the message with Beijing time
        const currentDate = new Date();
        const options = { timeZone: 'Asia/Shanghai', year: 'numeric', month: '2-digit', day: '2-digit', hour: '2-digit', minute: '2-digit', second: '2-digit', hour12: false };
        const timestamp = currentDate.toLocaleString('zh-CN', options).replace(/\//g, '-');
        const messageWithTimestamp = { ...message, timestamp, retainAfterPeek };

        let data = JSON.stringify(messageWithTimestamp);
        this.queues[taskID].push(data);

        // console.log(`enqueue taskID = ${taskID}, message = ${data}`);
    }

    // Remove and return the first message in the queue of the specific taskID
    dequeue(taskID) {
        if (!this.queues[taskID] || this.queues[taskID].length === 0) return null;

        const message = this.queues[taskID].shift();
        return { message, index: 0 };
    }

    // Consume all messages from the queue of the specific taskID
    dequeueAll(taskID) {
        if (!this.queues[taskID]) return [];

        const messages = this.queues[taskID].map((message, index) => ({ message, index }));
        this.queues[taskID] = [];
        return messages;
    }

    // Peek from specified position in the queue of the specific taskID
    peek(taskID, fromIndex, num) {
        // console.log(`peek fromIndex = ${fromIndex}, num = ${num}`);
        if (!this.queues[taskID] || fromIndex >= this.queues[taskID].length) return [];

        const messagesToPeek = num === -1 || num === undefined || num === null ? this.queues[taskID].slice(fromIndex) : this.queues[taskID].slice(fromIndex, fromIndex + num);
        // console.log(`peek messagesToPeek = `, messagesToPeek);
        const messages = messagesToPeek
            .map((message, index) => {
                if (message === null) return null; // 跳过已标记为 null 的消息

                const parsedMessage = JSON.parse(message);
                if (!parsedMessage.retainAfterPeek) {
                    this.queues[taskID][fromIndex + index] = null;
                }
                return {
                    message,
                    index: fromIndex + index
                };
            })
            .filter((message) => message !== null); // 过滤掉 null 消息

        return messages;
    }

    // Peek all messages from the queue of the specific taskID
    peekAll(taskID) {
        if (!this.queues[taskID]) return [];

        return this.peek(taskID, 0, -1);
    }

    // Get the size of the queue of the specific taskID
    size(taskID) {
        if (!this.queues[taskID]) return 0;

        return this.queues[taskID].length;
    }

    // Clear all queues (remove all taskID queues)
    clearAll() {
        this.queues = {};
    }

    // Get all taskIDs
    getAllTaskIDs() {
        return Object.keys(this.queues);
    }
}

class TaskQueue {
    constructor() {
        this.queue = [];
        this.tasks = {};
        this.processing = false;
        this.currentTaskId = null;
    }

    // 添加任务到队列
    add(taskId, task) {
        console.log('训练任务已经加入队列，任务ID:', taskId);
        this.tasks[taskId] = { status: 'queued', result: null, error: null, proc: null };
        this.queue.push({ taskId, task });
        this._processNext();
        return taskId;
    }

    // 停止任务
    async stop(taskId) {
        if (this.tasks[taskId] && this.tasks[taskId].proc) {
            console.log(`## Stop task ${taskId}, pid = `, this.tasks[taskId].proc.pid);
            try {
                let ret = await this._killProcess(this.tasks[taskId].proc.pid);
                if (ret) {
                    console.log(`## Task ${taskId} killed successfully.`);
                    this.tasks[taskId].status = 'stopped';
                    return true;
                } else {
                    console.error(`## Failed to stop task ${taskId}:`, err);
                    return false;
                }
            } catch (error) {
                console.error(`## Failed to stop task ${taskId}:`, error);
                return false;
            }
        } else {
            console.error(`## Task ${taskId} not found or has no associated process.`);
            return false;
        }
    }

    // 获取任务状态
    getTaskStatus(taskId) {
        return this.tasks[taskId] || null;
    }

    // 获取当前任务ID
    getCurrentTaskId() {
        return this.currentTaskId;
    }

    // 内部方法: 处理队列中的下一个任务
    async _processNext() {
        // console.log('this.processing = ', this.processing);
        if (this.processing || this.queue.length === 0) {
            return;
        }

        this.processing = true;
        const { taskId, task } = this.queue.shift();
        this.currentTaskId = taskId;
        // console.log(`taskId = ${taskId}, task = `, task);

        try {
            this.tasks[taskId].status = 'running';
            const result = await task(this._setProcess.bind(this, taskId));
            this.tasks[taskId].status = 'completed';
            this.tasks[taskId].result = result;
        } catch (error) {
            // console.log('task failed!, error = ', error);
            this.tasks[taskId].status = 'failed';
            this.tasks[taskId].error = error.message || error.toString();
        } finally {
            this.processing = false;
            this.currentTaskId = null;
            this._processNext();
        }
    }

    // 内部方法: 杀死进程，可以调用SIGKILL杀死进程
    async _killProcess(pid, signal = 'SIGTERM') {
        return new Promise((resolve, reject) => {
            treeKill(pid, signal, (err) => {
                if (err) {
                    reject(false);
                } else {
                    resolve(true);
                }
            });
        });
    }

    // 内部方法: 设定任务的子进程实例
    _setProcess(taskId, proc) {
        if (this.tasks[taskId]) {
            this.tasks[taskId].proc = proc;
        }
    }

    // 清理函数/关闭当前运行的任务
    async clear() {
        console.log(`## Close Task Queue Gracefully...`);
        if (this.currentTaskId) {
            const currentTask = this.tasks[this.currentTaskId];
            if (currentTask && currentTask.proc) {
                try {
                    await this._killProcess(currentTask.proc.pid);
                    currentTask.status = 'stopped';
                    console.log(`## Task ${this.currentTaskId} was stopped.`);
                } catch (error) {
                    console.error(`## Failed to stop the task ${this.currentTaskId}:`, error);
                }
                this.currentTaskId = null;
            }
        }
    }
}

module.exports = { MessageQueue, TaskQueue };
